/*
 * Copyright (C) 2008 by The Regents of the University of California
 * Redistribution of this file is permitted under the terms of the GNU
 * Public License (GPL).
 *
 * @author Junghoo "John" Cho <cho AT cs.ucla.edu>
 * @date 3/24/2008
 */
 
#include "BTreeIndex.h"
#include "BTreeNode.h"

using namespace std;

/*
 * Open the index file in read or write mode.
 * Under 'w' mode, the index file should be created if it does not exist.
 * @param indexname[IN] the name of the index file
 * @param mode[IN] 'r' for read, 'w' for write
 * @return error code. 0 if no error
 */
RC BTreeIndex::open(const string& indexname, char mode)
{
	RC rc;
	if (rc = pf.open(indexname, mode) < 0) return rc;

	// get rootPid and treeHeight
	if (rc = pf.read(0, treeInfo_buffer) < 0) return rc;
	memcpy(&rootPid, treeInfo_buffer, sizeof(rootPid));
	memcpy(&treeHeight, treeInfo_buffer + sizeof(rootPid), sizeof(treeHeight));

	if (!treeHeight)
		rootPid = -1;

	return 0;
}

/*
 * Close the index file.
 * @return error code. 0 if no error
 */
RC BTreeIndex::close()
{
    return pf.close();
}

/*
 * Insert (key, RecordId) pair to the index.
 * @param key[IN] the key for the value inserted into the index
 * @param rid[IN] the RecordId for the record being inserted into the index
 * @return error code. 0 if no error
 */
RC BTreeIndex::insert(int key, const RecordId& rid)
{
	RC rc;
	BTLeafNode leaf;

	if (rootPid == -1) { // we have an empty tree
		if (rc = leaf.insert(key, rid) < 0)	return rc;
		
		if (rc = leaf.write(1, pf) < 0) return rc;

		rootPid = 1;
		treeHeight = 1;

		// update tree info
		if (rc = updateTreeInfo() < 0) return rc;
	} else {
		if (treeHeight == 1) { // the root is also the leaf node
			leaf.read(1, pf);

			if (leaf.getKeyCount() == leaf.MAX_ENTRY) { // full node
				BTLeafNode sibling;
				int siblingKey;
				if (rc = leaf.insertAndSplit(key, rid, sibling, siblingKey) < 0) return rc;

				// update 2 nodes
				if (rc = leaf.write(1, pf) < 0)	return rc;
				if (rc = sibling.write(2, pf) < 0) return rc;
				if (rc = leaf.setNextNodePtr(2) < 0) return rc;

				// create a new node
				BTNonLeafNode root;
				if (rc = root.initializeRoot(1, siblingKey, 2) < 0) return rc;
				if (rc = root.write(3, pf) < 0) return rc;

				rootPid = 3;
				treeHeight = 2;
			} else {
				if (rc = leaf.insert(key, rid) < 0) return rc;
				if (rc = leaf.write(1, pf) < 0) return rc; 
			}

			// update treeInfo
			if (rc = updateTreeInfo() < 0) return rc;
		} else { // we have at least 2 levels
			insertHelper(-1, 1, rootPid, key, rid);
		}
	}

    return 0;
}


RC BTreeIndex::insertHelper(PageId parentNode, int currentLevel, PageId currentNode, const int &key, const RecordId &rid) {
	RC rc;

	if (currentLevel == treeHeight) { // leaf
		BTLeafNode leaf;
		if (rc = leaf.read(currentNode, pf) < 0) return rc;
		if (leaf.getKeyCount() == leaf.MAX_ENTRY) { // full node
			BTLeafNode sibling;
			int siblingKey;
			if (rc = leaf.insertAndSplit(key, rid, sibling, siblingKey) < 0) return rc;

			// update 2 nodes
			if (rc = leaf.write(currentNode, pf) < 0)	return rc;
			if (rc = sibling.write(pf.endPid() + 1, pf) < 0) return rc;
			if (rc = leaf.setNextNodePtr(pf.endPid() + 1) < 0) return rc;

			return siblingKey; 

		} else {
			if (rc = leaf.insert(key, rid) < 0) return rc;
			if (rc = leaf.write(currentNode, pf) < 0) return rc;
		}
	} else {
		// read the content of the current node
		BTNonLeafNode non_leaf;
		if (rc = non_leaf.read(currentNode, pf) < 0) return rc;
		
		// determine the child node
		PageId childNode;
		if (rc = non_leaf.locateChildPtr(key, childNode) < 0) return rc;
		
		// follow this child node to the leaf
		if (rc = insertHelper(currentNode, ++currentLevel, childNode, key, rid) <= 0) return rc;

		// rc now contains the siblingKey. We have to determine if the current non-leaf node is full
		// to push the key to the upper level
		int midKey = rc;
		if (rc = non_leaf.insert(midKey, childNode) < 0) {
			if (rc == RC_NODE_FULL) {
				int newMidKey;
				BTNonLeafNode nf_sibling;
				if (rc = non_leaf.insertAndSplit(midKey, childNode, nf_sibling, newMidKey) < 0) return rc;
				nf_sibling.write(pf.endPid() + 1, pf);
				
				// are we at the root node?
				if (currentLevel == 1) {
					// create a new root
					BTNonLeafNode newRoot;
					if (rc = newRoot.initializeRoot(currentNode, newMidKey, childNode) < 0) return rc;

					// write to disk
					PageId newRoot_pid = pf.endPid() + 1;
					if (rc = newRoot.write(newRoot_pid, pf) < 0) return rc;

					// update rootPid and treeHeight
					rootPid = newRoot_pid;
					++treeHeight;
					updateTreeInfo();

					return 0;
				} else {
					return newMidKey;
				}
			}
		}

	}


	return 0;
}

/*
 * Find the leaf-node index entry whose key value is larger than or 
 * equal to searchKey, and output the location of the entry in IndexCursor.
 * IndexCursor is a "pointer" to a B+tree leaf-node entry consisting of
 * the PageId of the node and the SlotID of the index entry.
 * Note that, for range queries, we need to scan the B+tree leaf nodes.
 * For example, if the query is "key > 1000", we should scan the leaf
 * nodes starting with the key value 1000. For this reason,
 * it is better to return the location of the leaf node entry 
 * for a given searchKey, instead of returning the RecordId
 * associated with the searchKey directly.
 * Once the location of the index entry is identified and returned 
 * from this function, you should call readForward() to retrieve the
 * actual (key, rid) pair from the index.
 * @param key[IN] the key to find.
 * @param cursor[OUT] the cursor pointing to the first index entry
 *                    with the key value.
 * @return error code. 0 if no error.
 */
RC BTreeIndex::locate(int searchKey, IndexCursor& cursor)
{
	if(treeHeight == 0) return RC_END_OF_TREE;
	if(treeHeight == 1) // if there is only 1 node in the B+ Tree, the root node is a leaf node.
	{
		BTLeafNode Root;  //create a leafnode object for root
		Root.read(rootPid,pf); // read the content in rootPid into Root.buffer
		Root.convert_buffer(); // convert Root.buffer into Root.buffer_int, it is necessary before calling Root.locate()
		RC rc;
		int Eid;
		if(rc = Root.locate(searchKey, Eid) == RC_NO_SUCH_RECORD) return RC_NO_SUCH_RECORD; 
		else {cursor.eid = Eid; cursor.pid = rootPid;}
		return 0;
	}
	else // if there are more than 1 node in the tree, the root node is a non-leaf node.
	{
		BTNonLeafNode Root;  //create a non-leaf node object for root
		Root.read(rootPid,pf);
		BTNonLeafNode * CurrentNode = &Root;
		int depth = 1;
		PageId next_pid;
		while (depth < treeHeight-1)
		{
			CurrentNode->locateChildPtr(searchKey,next_pid);  // get the pid of the proper child non-leaf node
			BTNonLeafNode NewNode;  // create a new node
			NewNode.read(next_pid,pf); // read the data from the next_pid page into NewNode.buffer;
			CurrentNode = &NewNode;  // let CurrenNode point at this new node;
			depth++;
		}
		// after the while loop, depth==treeHeight-1 and CurrentNode is pointing at the last nonleaf node NewNode.
		CurrentNode->locateChildPtr(searchKey,next_pid);  
		BTLeafNode leaf;  
		leaf.read(next_pid,pf);
		leaf.convert_buffer();
		RC rc;
		int Eid;
		if(rc = leaf.locate(searchKey, Eid) == RC_NO_SUCH_RECORD) return RC_NO_SUCH_RECORD; 
		else {cursor.eid = Eid; cursor.pid = rootPid;}
		return 0;
	}
}

/*
 * Read the (key, rid) pair at the location specified by the index cursor,
 * and move foward the cursor to the next entry.
 * @param cursor[IN/OUT] the cursor pointing to an leaf-node index entry in the b+tree
 * @param key[OUT] the key stored at the index cursor location.
 * @param rid[OUT] the RecordId stored at the index cursor location.
 * @return error code. 0 if no error
 */
RC BTreeIndex::readForward(IndexCursor& cursor, int& key, RecordId& rid)
{
	BTLeafNode fn;  // create a leaf node
	fn.read(cursor.pid,pf);  // read the data from page cursor.pid into fn.buffer
	fn.convert_buffer();   // convert fn.buffer into fn.buffer_int, it is necessary before calling fn.readEntry()
	RC rc;
	if(rc = fn.readEntry(cursor.eid, key, rid) != 0) return rc;  // if eid is larger than getKeyCount(), return RC_NO_SUCH_RECORD
	cursor.eid++;
    return 0;
}

/**
* Write rootPid & treeHeight to pagePid = 0
*/
RC BTreeIndex::updateTreeInfo() {
	// update rootPid and treeHeight
	memcpy(treeInfo_buffer, &rootPid, sizeof(rootPid));
	memcpy(treeInfo_buffer + sizeof(rootPid), &rootPid, sizeof(rootPid));
	return pf.write(0, treeInfo_buffer);
}
